{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# {class}`~tvm.relay.frontend.onnx.BatchNorm`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import set_env\n",
    "from tvm.relay.frontend.onnx import BatchNorm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0;31mInit signature:\u001b[0m \u001b[0mBatchNorm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mSource:\u001b[0m        \n",
      "\u001b[0;32mclass\u001b[0m \u001b[0mBatchNorm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mOnnxOpConverter\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0;34m\"\"\"Operator converter for BatchNorm.\"\"\"\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0;34m@\u001b[0m\u001b[0mclassmethod\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0;32mdef\u001b[0m \u001b[0m_impl_v1\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcls\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mattr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0;31m# TODO(zhreshold): 'spatial' is not properly handled here.\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0;31m# TODO(vvchernov): 'training_mode' (onnx tag) is not correctly handled, ignore for now\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mAttrCvt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m            \u001b[0mop_name\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"batch_norm\"\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m            \u001b[0mignores\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"spatial\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"is_test\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"consumed_inputs\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"momentum\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"training_mode\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mattr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0;31m# We only support test mode, so we return data, moving_mean, moving_var,\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0;31m# and then moving_mean and moving_var again as placeholders for\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0;31m# the expected \"saved_mean\", \"saved_var\".\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0;32mreturn\u001b[0m \u001b[0m_expr\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTupleWrapper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0m_expr\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTuple\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m5\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFile:\u001b[0m           /media/pc/data/lxw/ai/tvm/python/tvm/relay/frontend/onnx.py\n",
      "\u001b[0;31mType:\u001b[0m           type\n",
      "\u001b[0;31mSubclasses:\u001b[0m     "
     ]
    }
   ],
   "source": [
    "BatchNorm??"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0;31mInit signature:\u001b[0m \u001b[0m_expr\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTupleWrapper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtuple_value\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msize\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mDocstring:\u001b[0m     \n",
      "TupleWrapper.\n",
      "\n",
      "This class is a Python wrapper for a Relay tuple of known size.\n",
      "It allows for accessing the fields of the Relay tuple as though\n",
      "it were a Python tuple.\n",
      "\n",
      "Parameters\n",
      "----------\n",
      "tuple_value: tvm.relay.Expr\n",
      "    The input tuple\n",
      "\n",
      "size: int\n",
      "    The size of the tuple.\n",
      "\u001b[0;31mFile:\u001b[0m           /media/pc/data/lxw/ai/tvm/python/tvm/relay/expr.py\n",
      "\u001b[0;31mType:\u001b[0m           type\n",
      "\u001b[0;31mSubclasses:\u001b[0m     "
     ]
    }
   ],
   "source": [
    "from tvm.relay import expr as _expr\n",
    "_expr.TupleWrapper?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.relay.frontend.common import AttrCvt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = AttrCvt(\n",
    "    op_name=\"batch_norm\",\n",
    "    ignores=[\"spatial\", \"is_test\", \"consumed_inputs\", \"momentum\", \"training_mode\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "xin",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
